package com.ficus.road.maintain.util;

import com.ficus.road.maintain.core.model.Bbox;
import com.ficus.road.maintain.core.model.RelBbox;
import ficus.car_demo.FrameMeta;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.DataBuffer;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Base64;
import java.util.List;
import java.util.stream.IntStream;

@Slf4j
public class ImageUtil {
    public static byte[] convertRgb2Rgba(FrameMeta.RoadPerceptionResultSet roadPerceptionResultSet) throws IOException {

        byte[] bytes = roadPerceptionResultSet.getSegmentPicture().toByteArray();
        StringBuilder decodeStringBuilder = new StringBuilder();
        FrameMeta.ImageType pictureType = roadPerceptionResultSet.getPictureType();
        byte[] bytes1;
        if (FrameMeta.ImageType.IMAGE_TYPE_RLE_CODE_VALUE == pictureType.getNumber()) {
            FrameMeta.RunLengthEncoding runLengthEncoding = FrameMeta.RunLengthEncoding.parseFrom(bytes);
            List<FrameMeta.RunLengthEncoding.Segment> segmentListList = runLengthEncoding.getSegmentListList();
            int segmentHeight = runLengthEncoding.getHeight();
            int segmentWidth = runLengthEncoding.getWidth();
            for (FrameMeta.RunLengthEncoding.Segment s : segmentListList) {
                int length = s.getLength();
                for (int i = length; i > 0; i--) {
                    decodeStringBuilder.append((char) s.getLabel());
                }
            }
            bytes = decodeStringBuilder.toString().getBytes();
            BufferedImage bufferedImage = new BufferedImage(segmentWidth, segmentHeight, BufferedImage.TYPE_4BYTE_ABGR);
            for (int i = 0; i < segmentWidth; i++) {
                for (int j = 0; j < segmentHeight; j++) {
                    int elem = bytes[i + j * segmentWidth];
                    bufferedImage.setRGB(i, j, getAlphaByrGay(elem).getRGB());
                }
            }
            bytes1 = imageToBytes(bufferedImage);
        } else {
            long time = System.currentTimeMillis();
            InputStream buffIn = new ByteArrayInputStream(bytes, 0, bytes.length);
            BufferedImage bi = ImageIO.read(buffIn);
            if (bi == null) {
                return null;
            }
            int width = bi.getWidth();
            int height = bi.getHeight();
            BufferedImage bufferedImage = new BufferedImage(width, height, BufferedImage.TYPE_4BYTE_ABGR);
            DataBuffer dataBuffer = bi.getData().getDataBuffer();
            IntStream.rangeClosed(0, width - 1)
                    .parallel()
                    .forEach(i -> {
                        for (int j = 0; j < height; j++) {
                            int elem = dataBuffer.getElem(i + j * width);
                            bufferedImage.setRGB(i, j, getAlphaByrGay(elem).getRGB());
                        }
                    });
            bytes1 = imageToBytes(bufferedImage);
            log.debug("time={}", System.currentTimeMillis() - time);
        }
        return bytes1;
    }


    /**
     * BufferedImage转byte[]
     *
     * @param bImage BufferedImage对象
     * @return byte[]
     * @auth zhy
     */
    private static byte[] imageToBytes(BufferedImage bImage) {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        try {
            ImageIO.write(bImage, "png", out);
        } catch (IOException e) {
            log.error("IOException", e);
        }
        return out.toByteArray();
    }

    private static Color getAlphaByrGay(int gray) {
        switch (gray) {
            case 0:
                return new Color(0, 0, 0, 0);
            case 1:
                return new Color(120, 255, 114, 127);
            case 2:
                return new Color(120, 255, 114, 178);
            case 3:
                return new Color(120, 255, 114, 127);
            case 4:
                return new Color(105, 253, 255, 127);
            case 5:
                return new Color(105, 253, 255, 127);
            case 6:
                return new Color(255, 122, 0, 127);
            case 7:
                return new Color(105, 253, 255, 127);
            case 8:
                return new Color(255, 122, 0, 127);
            case 9:
                return new Color(255, 122, 0, 127);
            case 10:
                return new Color(117, 117, 117, 127);
            case 11:
                return new Color(105, 122, 0, 127);
        }
        return new Color(0, 0, 0, 0);
    }

    public static String bytes2Base64(byte[] bytes) {
        return Base64.getEncoder().encodeToString(bytes);
    }

    public static byte[] base64ToBytes(String base64) {
        return Base64.getDecoder().decode(base64);
    }

    public static byte[] drawFrame(String imageBase64, List<Bbox> bboxList, List<RelBbox> relBboxList) {
        Base64.Decoder decoder = Base64.getDecoder();
        byte[] originImageBytes = decoder.decode(imageBase64);
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ByteArrayInputStream bais = new ByteArrayInputStream(originImageBytes);
        try {
            BufferedImage image = ImageIO.read(bais);
            Graphics graphics = image.getGraphics();
            graphics.setColor(Color.RED);
            if (CollectionUtils.isNotEmpty(relBboxList)) {
                // 如果只有相对位置
                int width = image.getWidth();
                int height = image.getHeight();
                for (RelBbox relBbox : relBboxList) {
                    graphics.drawRect((int) (width * relBbox.getX()), (int) (height * relBbox.getY()), (int) (width * relBbox.getW()), (int) (height * relBbox.getH()));
                }
            } else if (CollectionUtils.isNotEmpty(bboxList)) {
                for (Bbox bbox : bboxList) {
                    graphics.drawRect(bbox.getX(), bbox.getY(), bbox.getW(), bbox.getH());
                }
            }
            graphics.dispose();
            ImageIO.write(image, "jpg", baos);
        } catch (IOException e) {
            log.error("decode image met exception:", e);
        } finally {
            try {
                bais.close();
            } catch (IOException e) {
                log.error("close oss stream met exception:", e);
            }
        }
        return baos.toByteArray();
    }
}
